import os
import subprocess

def download_realesrgan_models():
    """Downloads the pre-trained models for Real-ESRGAN."""
    print("--- Downloading Real-ESRGAN models ---")
    weights_dir = 'Real-ESRGAN/weights'
    os.makedirs(weights_dir, exist_ok=True)

    # RealESRGAN_x4plus.pth
    url_x4plus = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
    subprocess.run(['wget', url_x4plus, '-P', weights_dir])

    # RealESRGAN_x4plus_anime_6B.pth
    url_anime = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth'
    subprocess.run(['wget', url_anime, '-P', weights_dir])
    print("--- Real-ESRGAN models downloaded ---")

def download_rife_models():
    """Downloads the pre-trained models for RIFE."""
    print("--- Downloading RIFE models ---")
    train_log_dir = 'RIFE/train_log'
    os.makedirs(train_log_dir, exist_ok=True)

    # Download the model pack from Google Drive
    gdown_id = '1APIzVeI-4ZZCEuIRE1m6WYfSCaOsi_7_'
    zip_path = os.path.join(train_log_dir, 'RIFE_trained_model_v3.6.zip')
    subprocess.run(['gdown', '--id', gdown_id, '-O', zip_path])

    # Unzip the models
    subprocess.run(['7z', 'e', zip_path, f'-o{train_log_dir}'])
    print("--- RIFE models downloaded and extracted ---")

if __name__ == '__main__':
    download_realesrgan_models()
    download_rife_models()
    print("All models have been downloaded.")
